// func.cu
// 单元测试通过后再标记ok
// 修改或优化函数后一定要与ok的函数比对结果

#include <iostream>  // cout
#include <fstream>  // fstream
#include <vector>  // vector
#include "timer.h"  // 计时器
#include "cmdline.h"  // 解析器
#include "func.h"  // 数据结构与函数

#define HTD cudaMemcpyHostToDevice
#define DTH cudaMemcpyDeviceToHost

// init 初始化 ok
void init(int argc, char **argv, Option &option) {
  // 解析参数
  cmdline::parser parser;
  parser.add<std::string>("input", 'i', "input directory", true, "");
  parser.add<std::string>("output", 'o', "output directory", true, "");
  parser.add<int>("word", 'w', "word length 4-8",
    true, 4, cmdline::range(4, 8));
  parser.add<int>("precise", 'p', "precise mode 0 or 1",
    false, 0, cmdline::range(0, 1));
  parser.add<float>("similarity", 's', "similarity 0.1-1.0",
    true, 0.95, cmdline::range(0.1, 1.0));
  parser.parse_check(argc, argv);
  option.inputDir = parser.get<std::string>("input");  // 输入文件路径
  option.outputDir = parser.get<std::string>("output");  // 输出文件路径
  option.wordLength = parser.get<int>("word");  // 短词长度
  option.precise = parser.get<int>("precise");  // 精确模式
  option.similarity = parser.get<float>("similarity");  // 相似度阈值
  if (option.similarity < 0.8) option.precise = 1;  // 相似度低 必须精确模式
}

// readIndex 读数据索引 ok
void readIndex(Option &option, Index &index) {
  // 常用数据
  std::string line;  // 读一行数据
  int readsCount=0;  // 序列数
  {  // 读数据信息
    std::ifstream infoFile(option.inputDir+"/info.txt");  // 读配置文件
    getline(infoFile, line);  // 总序列数
    option.readsCount = std::stol(line);
    readsCount = option.readsCount;
    getline(infoFile, line);  // 最长序列
    option.longest = std::stol(line);
    getline(infoFile, line);  // 最短序列
    option.shortest = std::stol(line);
    infoFile.close();
    std::cout << "Reads count:\t" << option.readsCount << std::endl;
    std::cout << "Longest:\t" << option.longest << std::endl;
    std::cout << "Shortest:\t" << option.shortest << std::endl;
    std::cout << std::endl;  // 隔开下个信息块
  }
  {  // 读取索引信息
    index.nameIndex.resize(readsCount);  // 序列名索引
    index.readIndex.resize(readsCount);  // 序列数据索引
    index.readLengths.resize(readsCount);  // 序列长度索引
    std::ifstream indexFile(option.inputDir+"/index.bin");  // 序列索引文件
    indexFile.read((char*)index.nameIndex.data(), readsCount*sizeof(long));
    indexFile.read((char*)index.readIndex.data(), readsCount*sizeof(long));
    indexFile.read((char*)index.readLengths.data(), readsCount*sizeof(int));
    indexFile.close();
  }
}

// makeData 生成数据 ok
void makeData(char *read, int length, unsigned int *buf) {
  unsigned int *packed = buf+8;  // 打包数据
  unsigned int low = 0;  // 打包后的低位
  unsigned int high = 0;  // 打包后的高位
  int p = 0;  // 打包了多少碱基
  int a=0, c=0, g=0, t=0;  // 碱基数
  unsigned int base = 0;  // 碱基
  for (int i=0; i<length; i++, read++) {
    switch (*read) {  // 碱基转数字
      case 'a': base = 0; a+=1; break;
      case 'c': base = 1; c+=1; break;
      case 'g': base = 2; g+=1; break;
      case 't': base = 3; t+=1; break;
      case 'u': base = 3; t+=1; break;
      case 'A': base = 0; a+=1; break;
      case 'C': base = 1; c+=1; break;
      case 'G': base = 2; g+=1; break;
      case 'T': base = 3; t+=1; break;
      case 'U': base = 3; t+=1; break;
      default: base = 4; break;
    }
    if (base < 4) {  // 打包
      low >>= 1;
      low += (base&1)<<31;
      high >>= 1;
      high += (base>>1)<<31;
      p += 1;
      if (p == 32) {
        *packed = low;
        *(packed+1) = high;
        low = 0;
        high = 0;
        packed += 2;
        p = 0;
      }
    }
  }
  int netLength = a+c+g+t;
  low >>= (32-netLength%32);
  high >>= (32-netLength%32);
  *packed = low;
  *(packed+1) = high;
  *buf = length;
  *(buf+1) = netLength;
  *(buf+2) = a;
  *(buf+3) = c;
  *(buf+4) = g;
  *(buf+5) = t;
}

// readData 读数据 ok
void readData(Option &option, Index &index, Data &data) {
  // 常用数据
  int readsCount = option.readsCount;  // 序列数
  long bufLength = 0; // 数据总长度
  Timer timer;  // 计时器
  {  // 序列长度与偏移 ok
    timer.start();
    cudaMallocHost(&data.offsets_h, sizeof(long)*readsCount);  // 数据偏移h
    cudaMalloc(&data.offsets_d, sizeof(long)*readsCount);  // ~d
    for (int i=0; i<readsCount; i++) {  // 生成偏移
      data.offsets_h[i] = bufLength;
      int length = index.readLengths[i];  // 长度
      int len32 = (index.readLengths[i]+31)/32*32;  // 长度32补齐
      bufLength += 8+len32/32*2;  // 数据总长度
    }
    if (data.offsets_h[readsCount-1] == 0) {  // 序列数不整除线程
      data.offsets_h[readsCount-1] = bufLength;
      bufLength += 8;  // 防访存越界
    }
    cudaMallocHost(&data.buf_h, sizeof(int)*bufLength);  // 数据h
    cudaMalloc(&data.buf_d, sizeof(int)*bufLength);  // ~d
    cudaDeviceSynchronize();
    std::cout << "Compute offset:\t"; timer.getDuration();
  }
  {  // 读数据 ok
    timer.start();
    char buffer[65536];
    std::ifstream readFile(option.inputDir+"/read.txt");  // 序列文件
    for (int i=0; i<readsCount; i++) {  // 读序列数据
      readFile.seekg(index.readIndex[i], std::ios::beg);  // 跳转到数据头
      int length = index.readLengths[i];  // 序列长度
      readFile.read(buffer, length);  // 写入数据
      makeData(buffer, length, data.buf_h+data.offsets_h[i]);
    }
    readFile.close();
    cudaDeviceSynchronize();
    std::cout << "Read data:\t"; timer.getDuration();
  }
  // 收尾
  cudaMemcpy(data.buf_d, data.buf_h, sizeof(int)*bufLength, DTH);
  cudaMemcpy(data.offsets_d, data.offsets_h, sizeof(long)*readsCount, HTD);
  cudaDeviceSynchronize();
  std::cout << std::endl;  // 隔开下个信息块
}

// initBench 初始化bench ok
void initBench(Option &option, Bench &bench) {
  // 常用数据
  int readsCount = option.readsCount;  // 每线程处理序列数
  // 初始化
  bench.top = 0;
  cudaMallocHost(&bench.remains, sizeof(int)*readsCount);  // 剩余序列h
  for (int i=0; i<readsCount; i++) bench.remains[i] = i;  // 赋初值
  bench.remainCount = readsCount;  // 剩余序列数
  cudaMallocHost(&bench.jobs_h, sizeof(int)*readsCount*2);  // 任务h
  cudaMalloc(&bench.jobs_d, sizeof(int)*readsCount*2);  // ~d
  for (int i=0; i<readsCount; i++) {  // 初始化 (-1:失败 0:成功)
    bench.jobs_h[i*2] = i;
    bench.jobs_h[i*2+1] = -1;
  }
  bench.jobCount = readsCount;  // 任务数
  cudaMallocHost(&bench.cluster, sizeof(int)*readsCount);  // 聚类结果
  for (int i=0; i<readsCount; i++) bench.cluster[i] = -1;  // 赋初值
  cudaMallocManaged(&bench.list, sizeof(int)*2048);  // 短词列表
  cudaDeviceSynchronize();
}

// updateRepresent 更新代表序列 ok
void updateRepresent(Option &option, Data &data, Bench &bench) {
  // 常用数据
  {  // 更新聚类结果
    for (int i=0; i<bench.jobCount; i++) {  // 遍历任务列表
      int index = bench.jobs_h[i*2];
      int result = bench.jobs_h[i*2+1];
      if (result == 0) bench.cluster[index] = bench.top;  // 聚类成功
    }
  }
  {  // 更新 未聚类列表 任务列表
    int count = 0;  // 聚类未成功数
    for (int i=0; i<bench.remainCount; i++) {
      int index = bench.remains[i];
      if (bench.cluster[index] == -1) {  // 如果这个序列还没聚类
        bench.remains[count] = index;  // 未聚类列表
        bench.jobs_h[count*2] = index;  // 任务列表
        bench.jobs_h[count*2+1] = 0;  // 任务需要执行
        count += 1;
      }
    }
    bench.remainCount = count;
    bench.jobCount = count;
    cudaMemcpy(bench.jobs_d, bench.jobs_h, sizeof(int)*count*2, HTD);
  }
  {  // 更新代表序列
    if (bench.remainCount > 0) {  // 如果聚类还没完成
      bench.top = bench.remains[0];
    } else {
      bench.top = option.readsCount;  // 聚类完成
    }
  }
  {  // 广播代表序列
    if (bench.top == option.readsCount) return;  // 聚类完成
    if (bench.remainCount > 0) bench.cluster[bench.top] = bench.top;
  }
  cudaDeviceSynchronize();
}

// generateList 生成短词列表 ok
void generateList(Data &data, Option &option, Bench &bench) {
  unsigned int *list = bench.list;  // 短词列表
  int wordLength = option.wordLength;  // 短词长度
  memset(list, 0x00, sizeof(int)*2048);  // 清零 2048=65536/32
  unsigned int *bufBegin = data.buf_h+data.offsets_h[bench.top];  // 数据起始
  int netLength = *(bufBegin+1);  // 净长度
  unsigned int *packed = bufBegin+8;
  unsigned short mask = (1<<wordLength*2)-1;  // 短词的mask
  //　计算
  unsigned short word = 0;  // 短词
  for (int i=0, k=0; i<wordLength-1; i++, k++) {  // 短词初始化
    word <<= 1;
    word += *(packed+1)>>k&1;
    word <<= 1;
    word += *(packed)>>k&1;
  }
  for (int i=wordLength-1, k=wordLength-1; i<netLength; i++, k++) {
    word <<= 1;
    word += *(packed+1)>>k&1;
    word <<= 1;
    word += *(packed)>>k&1;
    word &= mask;
    int order = word/32;  // 第几个位
    unsigned int key = 1<<word%32;  // 滑动个数
    list[order] &= ~key;
    list[order] += key;
    if (k == 31) {k = -1; packed += 2;}
  }
  cudaDeviceSynchronize();
}

// updateJobs 更新任务列表 ok
void updateJobs(Bench &bench) {
  int count = 0;  // 剩余序列数
  int *result = bench.jobs_h;  // 上次的结果
  int *job = bench.jobs_h;  // 新的任务
  for (int i=0; i<bench.jobCount; i++, result+=2) {
    if (*(result+1) == 0) {  // -1:任务不通过 0:任务通过
      *job = *result;
      *(job+1) = 0;
      job += 2;
      count += 1;
    }
  }
  bench.jobCount = count;
  int byteCount = sizeof(int)*bench.jobCount*2;  // 要拷贝字节数
  cudaMemcpy(bench.jobs_d, bench.jobs_h, byteCount, HTD);
  cudaDeviceSynchronize();
}

// kernel_baseFilter 前置过滤 ok
__global__ void kernel_baseFilter(unsigned int *buf, long *offsets,
int top, int *jobs, int jobCount, float similarity) {
  int index = blockDim.x*blockIdx.x+threadIdx.x;  // 线程编号
  int loop = gridDim.x*blockDim.x;
  // 准备数据1
  unsigned int *represent = buf+offsets[top];
  int length1 = *represent;  // 长度
  int a1 = *(represent+2);  // 碱基数
  int c1 = *(represent+3);
  int g1 = *(represent+4);
  int t1 = *(represent+5);
  for (int idx=index; idx<jobCount; idx+=loop) {
    // 准备数据2
    unsigned int *bufBegin = buf+offsets[jobs[idx*2]];  // 起始位置
    int length2 = *bufBegin;
    int a2 = *(bufBegin+2);
    int c2 = *(bufBegin+3);
    int g2 = *(bufBegin+4);
    int t2 = *(bufBegin+5);
    // 计算
    int sum = 0;
    sum += min(a1, a2);
    sum += min(c1, c2);
    sum += min(g1, g2);
    sum += min(t1, t2);
    int length = min(length1, length2);
    length = ceil(((float)length)*similarity);
    if (sum < length) jobs[idx*2+1] = -1;  // 过滤失败
  }
}

// kernel_wordFilter 短词过滤 ok
__global__ void kernel_wordFilter(unsigned int *buf, long *offsets,
int top, int *jobs, int jobCount, int wordLength,
float similarity, unsigned int *listTemp) {
  // 准备数据
  __shared__ unsigned int list[2048];
  for (int i=threadIdx.x; i<2048; i+=WARP) {
    list[i] = listTemp[i];
  }
  int length1 = *(buf+offsets[top]);  // 长度
  unsigned short mask = (1<<wordLength*2)-1;  // 短词的mask
  int index = blockDim.x*blockIdx.x+threadIdx.x;  // 线程编号
  int loop = gridDim.x*blockDim.x;
  for (int idx=index; idx<jobCount; idx+=loop) {
    // 准备数据
    unsigned int *bufBegin = buf+offsets[jobs[idx*2]];  // 起始位置
    int length2 = *bufBegin;  // 长度
    int netLength2 = *(bufBegin+1);  // 净长度
    unsigned int *packed2 = bufBegin+8;  // 打包数据
    //　计算
    int sum = 0;
    unsigned short word = 0;  // 短词
    for (int i=0; i<netLength2; i+=32, packed2+=2) {
      for (int j=0; j<32; j++) {
        if (i+j == netLength2) break;
        word <<= 1;
        word += *(packed2+1)>>j&1;
        word <<= 1;
        word += *(packed2)>>j&1;
        word &= mask;
        int order = word/32;
        unsigned int key = 1<<word%32;
        if (i+j > wordLength-2) sum += ((list[order]&key) > 0);
      }
    }
    // 处理结果
    int length = min(length1, length2);
    int len = length-wordLength+1;  // 理论短词数
    length = ceil((float)length*(1.0f-similarity));
    int minLen = len-length*wordLength;
    if (sum < minLen) jobs[idx*2+1] = -1;  // 过滤失败
    __syncthreads();  // 同步一下 否则对性能有影响
  }
}

// kernel_dynamic 动态规划 ok
__global__ void kernel_dynamic(unsigned int *buf, long *offsets,
int top, int *jobs, int jobCount, float similarity) {
  // 准备数据1
  __shared__ unsigned int packed1[4096];  // 65536/32*2
  unsigned int *represent = buf+offsets[top];
  int length1 = *represent;  // 长度
  int netLength1 = *(represent+1);  // 净长度
  int netLen321 = (netLength1+31)/32*32;  // 净长度32补齐个数
  unsigned int *read = represent+8;
  for (int i=threadIdx.x; i<netLen321/32*2; i+=blockDim.x) {
    packed1[i] = read[i];
  }
  int index = blockDim.x*blockIdx.x+threadIdx.x;  // 线程编号
  int loop = gridDim.x*blockDim.x;
  for (int idx=index; idx<jobCount; idx+=loop) {
    // 准备数据2 指针偏移和下标的速度一样
    unsigned int *bufBegin = buf+offsets[jobs[idx*2]];  // 起始位置
    int length2 = *bufBegin;  // 长度
    int netLength2 = *(bufBegin+1);  // 净长度
    int netLen322 = (netLength2+31)/32*32;  // 净32补齐
    unsigned int *packed2 = bufBegin+8;  // 打包数据
    // 准备中间数据 1是行2是列
    unsigned int line[2048];  // 保存结果 二进制 2048=65536/32
    memset(line, 0xFF, 2048*sizeof(unsigned int));  // 0:匹配 1:不匹配
    int shift = ceil((float)length1-(float)length2*similarity);  // 偏移量
    shift = (shift+31)/32;  // 偏移块数
    int cols = 32;  // 竖向计算行数
    // 计算
    for (int i=0; i<netLen322/32; i++) {  // 遍历列
      if ((i+1)*32 > netLength2) cols = netLength2%32;
      unsigned int column[32] = {0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,
                                0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0};  // 进位
      unsigned int Low2 = packed2[i*2+0];  // 低32位
      unsigned int High2 = packed2[i*2+1];  // 高32位
      int jstart = max(i-shift, 0);
      int jend = min(i+shift, netLen321/32-1);
      for (int j=jstart; j<=jend; j++) {  // 遍历行
        unsigned int Low1 = packed1[j*2+0];  // 低32位
        unsigned int High1 = packed1[j*2+1];  // 高32位
        unsigned int row = line[j];  // 上一行结果
        for (int k=0; k<cols; k++) {  // 32*32的核心
          unsigned int low2 = 0x00000000;  // 1位扩展成32位
          if (Low2>>k&1) low2 = 0xFFFFFFFF;
          unsigned int high2 = 0x00000000;
          if (High2>>k&1) high2 = 0xFFFFFFFF;
          unsigned int temp1 = Low1 ^ low2;
          unsigned int temp2 = High1 ^ high2;
          unsigned int match = (~temp1)&(~temp2);
          unsigned int unmatch = ~match;
          unsigned int temp3 = row & match;
          unsigned int temp4 = row & unmatch;
          unsigned int carry = column[k];
          unsigned int temp5 = row + carry;
          unsigned int carry1 = temp5 < row;
          temp5 += temp3;
          unsigned int carry2 = temp5 < temp3;
          carry = carry1 | carry2;
          row = temp5 | temp4;
          column[k] = carry;
        }
        line[j] = row;
      }
    }
    // 统计结果
    int sum = 0;  // 比对得分
    int rows = 32;  // 横向计算列数
    for (int i=0; i<netLen321/32; i++) {
      if ((i+1)*32 > netLength1) rows = netLength1%32;
      unsigned int row = line[i];
      for (int j=0; j<rows; j++) sum += row>>j&1^1;
    }
    int cutoff = ceil((float)length2*similarity);
    if (sum < cutoff) jobs[idx*2+1] = -1;
  }
}

// clustering 聚类
void clustering(Option &option, Data &data, Bench &bench) {
  // 常用数据
  // int rank = option.rank;  // MPI进程号
  int wordLength = option.wordLength;  // 短词长度
  int precise = option.precise;  // 是否启用精确模式
  float similarity = option.similarity;  // 相似度阈值
  int readsCount = option.readsCount;  // 序列数
  initBench(option, bench);  // 初始化bench
  while (true) {  // 聚类
    {  // 准备工作
      std::cout << "\r" << bench.top << "/" << readsCount << std::flush;
      if (bench.top == readsCount) std::cout << "\n";
      // }
      if (bench.top == readsCount) break;  // 聚类完成
      generateList(data, option, bench);  // 生成短词列表
    }
    {  // 前置过滤 ok
      updateJobs(bench);
      kernel_baseFilter<<<BLOCK, WARP>>>
      (data.buf_d, data.offsets_d, bench.top,
      bench.jobs_d, bench.jobCount, similarity);  // 前置过滤
      int byteCount = sizeof(int)*bench.jobCount*2;  // 要拷贝字节数
      cudaMemcpy(bench.jobs_h, bench.jobs_d, byteCount, DTH);
    }
    if (!precise) {  // 短词过滤 ok
      updateJobs(bench);
      kernel_wordFilter<<<BLOCK, WARP>>>
      (data.buf_d, data.offsets_d, bench.top, bench.jobs_d, 
      bench.jobCount, wordLength, similarity, bench.list);  // 短词过滤
      int byteCount = sizeof(int)*bench.jobCount*2;  // 要拷贝字节数
      cudaMemcpy(bench.jobs_h, bench.jobs_d, byteCount, DTH);
    }
    {  // 序列比对 ok
      updateJobs(bench);
      kernel_dynamic<<<BLOCK, WARP>>>
      (data.buf_d, data.offsets_d, bench.top,
      bench.jobs_d, bench.jobCount, similarity);  // 动态规划
      int byteCount = sizeof(int)*bench.jobCount*2;  // 要拷贝字节数
      cudaMemcpy(bench.jobs_h, bench.jobs_d, byteCount, DTH);
    }
    updateRepresent(option, data, bench);  // 更新代表序列
  }
  std::cout << std::endl;  // 隔开下个信息块
}

// saveResult 保存结果 ok
void saveResult(Option &option, Index &index, Bench &bench) {
  // 常用数据
  int readsCount = option.readsCount;  // 序列数
  std::vector<std::vector<int>> represents;  // 代表序列
  std::string line;  // 一行数据
  int count = 0;  // 代表序列个数
  Timer timer; timer.start();  // 计时器
  {  // 找代表序列
    represents.resize(readsCount);
    for (int i=0; i<readsCount; i++) {  // 遍历结
      int top = bench.cluster[i];  // 这条序列所属的类
      represents[top].push_back(i);  // 这个类成员加一
    }
  }
  {  // 写入结果
    std::ifstream nameFile(option.inputDir+"/name.txt");  // 序列名文件
    std::ifstream readFile(option.inputDir+"/read.txt");  // 序列文件
    std::ofstream representFile(option.outputDir+"/represent.txt");  // 代表序列
    std::ofstream clusterFile(option.outputDir+"/cluster.fasta");  // 聚类结果
    for (int i=0; i<readsCount; i++) {  // 遍历代表序列
      if (represents[i].size() > 0) {  // 找到一个代表序列
        count += 1;  // 代表序列加一
        int order = represents[i][0];  // 代表序列编号
        nameFile.seekg(index.nameIndex[order], std::ios::beg);  // 跳转
        std::getline(nameFile, line);  // 读序列名
        clusterFile << line << "\n";  // 聚类结果
        representFile << line << "\n";  //代表序列
        for (int j=0; j<represents[i].size(); j++) {
          nameFile.seekg(index.nameIndex[represents[i][j]], std::ios::beg);
          std::getline(nameFile, line);  // 读序列名
          clusterFile << "  " << line << "\n";  // 类中的其他序列
        }
        readFile.seekg(index.readIndex[order], std::ios::beg);  // 跳转
        std::getline(readFile, line);  // 读序列
        representFile << line << "\n";  // 代表序列
      }
    }
    readFile.close();
    nameFile.close();
    representFile.close();
    clusterFile.close();
  }
  std::cout << "Cluster:\t" << count << std::endl;
  std::cout << "Save result:\t"; timer.getDuration();
}

// finish 收尾 ok
void finish(Data &data, Bench &bench) {
  cudaFreeHost(data.buf_h);  // 释放空间
  cudaFree(data.buf_d);
  cudaFreeHost(data.offsets_h);
  cudaFree(data.offsets_d);
  cudaFreeHost(bench.remains);
  cudaFreeHost(bench.jobs_h);
  cudaFree(bench.jobs_d);
  cudaFreeHost(bench.cluster);
  cudaFree(bench.list);
  cudaDeviceSynchronize();  // 同步
}
